# coding=utf-8
from numpy import *


def get_words():
    posting_list = [['my', 'dog', 'has', 'flea', 'problems', 'help', 'please'],
                    ['maybe', 'not', 'take', 'him', 'to', 'dog', 'park', 'stupid'],
                    ['my', 'dalmation', 'is', 'so', 'cute', 'I', 'love', 'him'],
                    ['stop', 'posting', 'stupid', 'worthless', 'garbage'],
                    ['mr', 'licks', 'ate', 'my', 'steak', 'how', 'to', 'stop', 'him'],
                    ['quit', 'buying', 'worthless', 'dog', 'food', 'stupid']]
    class_vec = [0, 1, 0, 1, 0, 1]
    return posting_list, class_vec


# create a vocabulary list, this include all words from data set.
# but make words distinct.
def create_vocab_list(data_set):
    voc_set = set([])
    for doc in data_set:
        voc_set = voc_set | set(doc)
    return list(voc_set)


# convert words list to a vector.
# the input set is one document from data set.
# function return a vector, show if the words from vocabulary occur in document.
# If occur, the val is 1, if not, val is 0.
def words2vec(vocab_list, input_set):
    vec = [0] * len(vocab_list)
    for word in input_set:
        if word in vocab_list:
            vec[vocab_list.index(word)] += 1
    return vec


# Calculate the probability that each of the word from vocabulary appear
# in each category.
# for example, word 'cute' appear one time in category 'normal',
# but there are 24 words in category 'normal' totally.
# So the p0 of word 'cute' is 1 / 24 = 0.4166667.
def train_nb0(train_mat, train_category):
    train_doc_num = len(train_mat)
    words_num = len(train_mat[0])

    p_abusive = sum(train_category) / float(train_doc_num)

    p0_num = ones(words_num)
    p1_num = ones(words_num)

    p0_denom = 2.0
    p1_denom = 2.0

    for i in range(train_doc_num):
        if train_category[i] == 1:
            p1_num += train_mat[i]
            p1_denom += sum(train_mat[i])
        else:
            p0_num += train_mat[i]
            p0_denom += sum(train_mat[i])

    p1_vec = log(p1_num / p1_denom)
    p0_vec = log(p0_num / p0_denom)

    return p0_vec, p1_vec, p_abusive


def classify_nb(vec_classify, vec_p0, vec_p1, class_p):
    p1 = sum(vec_classify * vec_p1) + log(class_p)
    p0 = sum(vec_classify * vec_p0) + log(1.0 - class_p)
    if p1 > p0:
        return 1
    else:
        return 0


def test_nb():
    posts, classes = get_words()
    vocab_list = create_vocab_list(posts)

    train_mat = []
    for doc in posts:
        train_mat.append(words2vec(vocab_list, doc))
    p0, p1, pab = train_nb0(array(train_mat), array(classes))

    test_sent = ["I", "love", "my", "dalmation"]
    doc = array(words2vec(vocab_list, test_sent))
    print("test '%s', classify as: %d" % (' '.join(test_sent), classify_nb(doc, p0, p1, pab)))

    test_sent = ["stupid", "garbage"]
    doc = array(words2vec(vocab_list, test_sent))
    print("test '%s', classify as: %d" % (' '.join(test_sent), classify_nb(doc, p0, p1, pab)))
